import os
import numpy as np

# Define the path where the robust model will be saved and loaded from.
ROBUST_MODEL_PATH = "lenet_robust_model.pth"

# --- Stage 1: Create the Robust Model via Adversarial Training ---
print("="*80)
print("Stage 1: Creating a robust model via adversarial fine-tuning...")
print(f"This will be saved to: {ROBUST_MODEL_PATH}")
print("="*80)

# This command corresponds to the "fine-tuning" bar in the paper's figure.
# We start from the pretrained clean model, enable training, and fine-tune all parameters.
train_command = (
    f"python mnist_lenet.py "
    f"--train_mode 1 "      # Set to training mode
    f"--adv_train "         # Enable adversarial training
    f"--paradigm 3 "        # Fine-tune all parameters
    f"--pretrain 1 "        # Start from the clean pretrained model
    f"--epsilon 0.3 "       # Use the standard epsilon for training
    f"--lr 0.0001 "         # A smaller learning rate is typical for fine-tuning
    f"--model_save_path {ROBUST_MODEL_PATH}"
)

# Execute the training command
print(f"\nExecuting training command:\n{train_command}\n")
os.system(train_command)

print("\n" + "="*80)
print("Stage 1 Finished. Robust model should be created.")
print("="*80 + "\n")


# --- Stage 2: Test the ROBUST Model with fixed lambda values ---
print("="*80)
print("Stage 2: Testing the robust model with different fixed lambda values...")
print(f"Loading model from: {ROBUST_MODEL_PATH}")
print("="*80)

# Check if the robust model was created successfully
if not os.path.exists(ROBUST_MODEL_PATH):
    print(f"Error: Robust model '{ROBUST_MODEL_PATH}' not found. Stage 1 might have failed.")
else:
    # Loop through lambda values from 0.0 to 1.0
    for lmbd_val in np.arange(0.0, 1.1, 0.1):
        lmbd = round(lmbd_val, 1)
        
        # This command loads the ROBUST model and tests it in --train_mode 0
        test_command = (
            f"python mnist_lenet.py "
            f"--train_mode 0 " # Set to testing mode
            f"--model_load_path {ROBUST_MODEL_PATH} " # CRITICAL: Load the robust model
            f"--lmbd {lmbd} "
            f"--no-console-log"
        )
        
        print("\n" + "-"*60)
        print(f"Executing test command: {test_command}")
        print("-"*60 + "\n")
        
        # Execute the test command
        os.system(test_command)

    print("\n--- All experiments for Figure 3 finished. ---") 